package com.example.websocket.server;

import org.java_websocket.WebSocket;
import org.java_websocket.handshake.ClientHandshake;
import org.java_websocket.server.WebSocketServer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.StringUtils;

import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/*
 * @className: WebSocketServer
 * @author: chentao
 * @date: 2025/10/24 上午9:22
 * @Version: 1.0
 * @description:
 */
public class SysWebSocketServer extends WebSocketServer {

    Logger logger = LoggerFactory.getLogger(SysWebSocketServer.class);

    private Map<String, WebSocket> socketMap = new ConcurrentHashMap<>();

    public static final String PREF = "/";

    public static final String SPACE = "";

    private int port;

    private WebSocketServerListener webSocketServerListener = new WebSocketServerListener() {
        @Override
        public void onOpen(WebSocket webSocket, ClientHandshake clientHandshake) {
            WebSocketServerListener.super.onOpen(webSocket, clientHandshake);
        }

        @Override
        public void onClose(WebSocket webSocket, int code, String reason, boolean remote) {
            WebSocketServerListener.super.onClose(webSocket, code, reason, remote);
        }

        @Override
        public void onError(WebSocket webSocket, Exception e) {
            WebSocketServerListener.super.onError(webSocket, e);
        }

        @Override
        public void onMessage(WebSocket webSocket, String data) {
            WebSocketServerListener.super.onMessage(webSocket, data);
        }

        @Override
        public void onMessage(WebSocket webSocket, ByteBuffer message) {
            WebSocketServerListener.super.onMessage(webSocket, message);
        }
    };

    public SysWebSocketServer(int port) {
        super(new InetSocketAddress(port));
        this.port = port;
    }

    public SysWebSocketServer(int port, WebSocketServerListener webSocketServerListener) {
        super(new InetSocketAddress(port));
        this.port = port;
        this.webSocketServerListener = webSocketServerListener;
    }

    @Override
    public void onOpen(WebSocket webSocket, ClientHandshake clientHandshake) {
        if (StringUtils.hasText(clientHandshake.getResourceDescriptor()) && clientHandshake.getResourceDescriptor().startsWith(PREF)) {
            logger.warn("websocket client serverId " + getServerKey(webSocket.getResourceDescriptor()) + " connect server success ...");
            this.socketMap.put(getServerKey(webSocket.getResourceDescriptor()), webSocket);
        }
        webSocketServerListener.onOpen(webSocket, clientHandshake);
    }

    @Override
    public void onClose(WebSocket webSocket, int i, String s, boolean b) {
        logger.warn("websocket client serverId " + getServerKey(webSocket.getResourceDescriptor()) + " down ...");
        this.socketMap.remove(getServerKey(webSocket.getResourceDescriptor()));
        webSocketServerListener.onClose(webSocket, i, s, b);
    }

    @Override
    public void onMessage(WebSocket webSocket, String data) {
        webSocketServerListener.onMessage(webSocket, data);
    }

    @Override
    public void onMessage(WebSocket webSocket, ByteBuffer message) {
        webSocketServerListener.onMessage(webSocket, message);
    }

    @Override
    public void onError(WebSocket webSocket, Exception e) {
        logger.warn("websocket client serverId " + getServerKey(webSocket.getResourceDescriptor()) + " down ...");
        this.socketMap.remove(getServerKey(webSocket.getResourceDescriptor()));
        webSocketServerListener.onError(webSocket, e);
    }

    @Override
    public void onStart() {
        logger.warn("websocket server started on port: " + this.port);
    }

    public String getServerKey(String resourceDescriptor) {
        if (resourceDescriptor.startsWith(PREF)) {
            return resourceDescriptor.replaceAll(PREF, SPACE);
        }
        return resourceDescriptor;
    }

    public long getActiveSize() {
        return this.socketMap.size();
    }

    public Set<String> getServerKeys() {
        return this.socketMap.keySet();
    }

    public WebSocket getWebSocket(String serverKey) {
        return this.socketMap.get(serverKey);
    }

    public void send(String serverKey, byte[] message) throws Exception {
        if (this.socketMap.containsKey(serverKey)) {
            this.socketMap.get(serverKey).send(message);
        } else {
            throw new Exception("websocket client serverId " + serverKey + " not connect server ...");
        }
    }

    public void sendAll(byte[] message) {
        broadcast(message);
    }

    public interface WebSocketServerListener {

        default public void onOpen(WebSocket webSocket, ClientHandshake clientHandshake) {

        }

        default public void onError(WebSocket webSocket, Exception e) {

        }

        default public void onClose(WebSocket webSocket, int code, String reason, boolean remote) {

        }

        default public void onMessage(WebSocket webSocket, String data) {

        }

        default public void onMessage(WebSocket webSocket, ByteBuffer message) {

        }
    }
}

